#include "Avalon/Services/ServiceNode.hpp"
#include <boost/bind.hpp>
#include "Avalon/Base/Algorithm.hpp"
#include "Avalon/IO/Reader.hpp"
#include "Avalon/Serialization/DataShuttleFactory.hpp"
#include "Avalon/Services/Service.hpp"
#include "Avalon/Services/TupleMessage.hpp"

using namespace Avalon;
using namespace Avalon::IO;
using namespace Avalon::Serialization;
using namespace Avalon::Services;
using namespace Avalon::Threading;
using namespace boost;
using namespace boost::posix_time;
using namespace boost::signals2;
using namespace std;

AVALON_DECLARE_MESSAGE(HeartbeatMessage, "Avalon.Services.HeartbeatMessage");

const time_duration ServiceNode::TIMEOUT = seconds(3);

ServiceNode::ServiceNode(DataShuttleFactory* dataShuttleFactory,
    time_duration timeout)
    : m_dataShuttleFactory(dataShuttleFactory),
      m_nextRequestId(0) {
  if(timeout != pos_infin) {
    m_heartbeatTimer.ConnectExpiredSignal(
      m_signalHandler.GetSlot<Timer::ExpiredSignal>(
      bind(&ServiceNode::OnHeartbeat, this)));
    m_heartbeatTimer.Start(timeout);
  }
  m_signalHandler.Start();
}

ServiceNode::~ServiceNode() {
  assert(m_channelEntries.empty());
  Clear(m_signals);
}

void ServiceNode::SendMessage(Channel* channel, const Message& message) {
  boost::lock_guard<boost::mutex> lock(m_mutex);
  map<Channel*, ChannelEntry*>::iterator channelIterator =
    m_channelEntries.find(channel);
  assert(channelIterator != m_channelEntries.end());
  ChannelEntry* entry = channelIterator->second;
  if(!entry->m_state.IsOpen()) {
    return;
  }
  entry->m_serializer->Shuttle("message", &const_cast<Message&>(message));
  channel->GetWriter().Write(entry->m_serializer->GetSerializedData());
}

connection ServiceNode::AddChannel(Channel* channel,
    const Connection::ClosedSignal::slot_type& slot, bool* isConnected) {
  boost::lock_guard<boost::mutex> lock(m_mutex);
  assert(m_channelEntries.find(channel) == m_channelEntries.end());
  ChannelEntry* entry = new ChannelEntry();
  m_channelEntries.insert(make_pair(channel, entry));
  entry->m_channel = channel;
  entry->m_serializer.reset(m_dataShuttleFactory->CreateSerializer());
  entry->m_deserializer.reset(m_dataShuttleFactory->CreateDeserializer());
  entry->m_sentHeartbeat = false;
  entry->m_timeoutCounter = 3;
  entry->m_isReadPending = false;
  bool isOpen;
  connection closedConnection = channel->GetConnection().ConnectClosedSignal(
    m_signalHandler.GetSlot<Connection::ClosedSignal>(
    bind(&ServiceNode::OnClosed, this, entry)), &isOpen);
  entry->m_state.Initialize(isOpen, closedConnection);
  *isConnected = isOpen;
  if(isOpen) {
    try {
      Async<int>::Ptr readResult = channel->GetReader().Read(&entry->m_data);
      entry->m_isReadPending = true;
      readResult->ConnectFinishedSignal(
        m_signalHandler.GetSlot<BaseAsync::FinishedSignal>(
        bind(&ServiceNode::OnRead, this, readResult, entry)));
    } catch(IOException&) {}
  }
  return entry->m_closedSignal.connect(slot);
}

void ServiceNode::RemoveChannel(Channel* channel) {
  boost::lock_guard<boost::mutex> lock(m_mutex);
  map<Channel*, ChannelEntry*>::iterator channelIterator =
    m_channelEntries.find(channel);
  assert(channelIterator != m_channelEntries.end());
  InternalRemoveChannel(channelIterator);
}

void ServiceNode::CloseAndRemoveChannel(Channel* channel) {
  boost::unique_lock<boost::mutex> lock(m_mutex);
  map<Channel*, ChannelEntry*>::iterator channelIterator =
    m_channelEntries.find(channel);
  assert(channelIterator != m_channelEntries.end());
  ChannelEntry* entry = channelIterator->second;
  if(entry->m_state.IsOpen()) {
    channel->GetConnection().Close();
    entry->m_closedCondition.wait(lock);
  }
  InternalRemoveChannel(channelIterator);
}

void ServiceNode::InternalRemoveChannel(
    map<Channel*, ChannelEntry*>::iterator& channelIterator) {
  ChannelEntry* entry = channelIterator->second;
  assert(!entry->m_state.IsOpen());
  m_channelEntries.erase(channelIterator);
  for(vector<BaseAsyncOwnerPtr>::const_iterator i = entry->m_responses.begin();
      i != entry->m_responses.end(); ++i) {
    (*i)->SetException(copy_exception(IOException("Channel removed.")));
  }
  delete entry;
}

void ServiceNode::OnClosed(ChannelEntry* entry) {
  boost::lock_guard<boost::mutex> lock(m_mutex);
  for(vector<BaseAsyncOwnerPtr>::const_iterator i = entry->m_responses.begin();
      i != entry->m_responses.end(); ++i) {
    (*i)->SetException(copy_exception(IOException("Connection closed.")));
  }
  entry->m_responses.clear();
  entry->m_state.SetClosed();
  if(!entry->m_isReadPending) {
    entry->m_closedCondition.notify_one();
    entry->m_closedSignal();
  }
}

void ServiceNode::OnRead(const Async<int>::Ptr& result, ChannelEntry* entry) {
  boost::lock_guard<boost::mutex> lock(m_mutex);
  entry->m_isReadPending = false;
  if(!entry->m_state.IsOpen()) {
    entry->m_closedCondition.notify_one();
    entry->m_closedSignal();
    return;
  }
  Channel* channel = entry->m_channel;
  if(result->GetState() == BaseAsync::EXCEPTION) {
    channel->GetConnection().Close();
    return;
  }
  entry->m_timeoutCounter = 3;
  scoped_ptr<Message> message;
  try {
    entry->m_deserializer->SetDataSource(entry->m_data.GetData(),
      result->Get());
    entry->m_deserializer->Shuttle("message", message);
    if(message == NULL) {
      channel->GetConnection().Close();
      return;
    }
  } catch(std::exception&) {
    channel->GetConnection().Close();
    return;
  }
  try {
    Async<int>::Ptr nextReadResult = channel->GetReader().Read(&entry->m_data);
    entry->m_isReadPending = true;
    nextReadResult->ConnectFinishedSignal(
      m_signalHandler.GetSlot<BaseAsync::FinishedSignal>(
      bind(&ServiceNode::OnRead, this, nextReadResult, entry)));
  } catch(IOException&) {}
  if(dynamic_cast<HeartbeatMessage*>(message.get()) != NULL) {
    return;
  }
  ServiceMessage* serviceMessage = dynamic_cast<ServiceMessage*>(
    message.get());
  if(serviceMessage != NULL && serviceMessage->IsResponseMessage()) {
    map<int, BaseAsyncPtr>::iterator responseIterator = m_pendingRequests.find(
      serviceMessage->GetRequestId());
    if(responseIterator != m_pendingRequests.end()) {
      BaseAsyncPtr result = responseIterator->second;
      m_pendingRequests.erase(responseIterator);
      Remove(entry->m_responses, dynamic_pointer_cast<BaseAsyncOwner>(result));
      serviceMessage->SetAsync(result);
    }
    return;
  }
  map<const type_info*, signal_base*>::iterator signalIterator =
    m_signals.find(&typeid(*message.get()));
  if(signalIterator == m_signals.end()) {
    return;
  }
  message->EmitSignal(signalIterator->second, this, channel);
}

void ServiceNode::OnHeartbeat() {
  boost::lock_guard<boost::mutex> lock(m_mutex);
  HeartbeatMessage message;
  for(map<Channel*, ChannelEntry*>::iterator i = m_channelEntries.begin();
      i != m_channelEntries.end(); ++i) {
    Channel* channel = i->first;
    ChannelEntry* channelEntry = i->second;
    --channelEntry->m_timeoutCounter;
    if(channelEntry->m_timeoutCounter != 0) {
      channelEntry->m_serializer->Shuttle("message", &message);
      channel->GetWriter().Write(
        channelEntry->m_serializer->GetSerializedData());
    } else {
      channel->GetConnection().Close();
    }
  }
}
